#!/usr/bin/python -tt

import parse
import time
import sys, os
import psycopg2
import psycopg2.extensions
import logging

# Log everything, and send it to stderr.
logging.basicConfig(level=logging.DEBUG)

def write_tag(count, type, name):
    logging.debug('Tag %s - %s: %s' % (count, type, name))
    cur.execute('SELECT upsert_tag(%s, %s, %s);', (name, type, str(count)))
    rslt = cur.fetchone()
    if rslt[0] == True:
        logging.debug('Updated tag ' + name)
        #raise Exception('DEBUG')


def print_tag(count, type, name):
    logging.debug('Tag %s - %s: %s' % (count, type, name))

def init_db():
    conn = psycopg2.connect(database='isidabot', user='postgres', port='5432')
    cur = conn.cursor()
    cur.execute('''
CREATE OR REPLACE FUNCTION make_plpgsql()
RETURNS VOID
LANGUAGE SQL
AS $$
CREATE LANGUAGE plpgsql;
$$;
 
SELECT
    CASE
    WHEN EXISTS(
        SELECT 1
        FROM pg_catalog.pg_language
        WHERE lanname='plpgsql'
    )
    THEN NULL
    ELSE make_plpgsql() END;
 
DROP FUNCTION make_plpgsql();
    ''')
    cur.execute('''
CREATE OR REPLACE FUNCTION create_tags_table ()
  RETURNS void AS $_$
BEGIN

IF EXISTS (
    SELECT *
    FROM   pg_catalog.pg_tables 
    WHERE  schemaname = 'public'
    AND    tablename  = 'tags'
    ) THEN
   RAISE NOTICE 'Table "public"."tags" already exists.';
ELSE
    CREATE TABLE public.tags (
      name  varchar(70) PRIMARY KEY,
      type varchar(14),
      count integer DEFAULT 0
    );
END IF;

END;
$_$ LANGUAGE plpgsql;
SELECT create_tags_table();
DROP FUNCTION create_tags_table();
    ''')
    cur.execute('ALTER TABLE tags OWNER TO "isidabot";')
    cur.execute('''
CREATE OR REPLACE FUNCTION upsert_tag(key TEXT, Atype TEXT, Acount INTEGER) RETURNS boolean AS $$
BEGIN
    LOOP
        -- first try to update the key
        UPDATE tags SET type = Atype, count = Acount WHERE name = key;
        IF found THEN
            RETURN TRUE;
        END IF;
        -- not there, so try to insert the key
        -- if someone else inserts the same key concurrently,
        -- we could get a unique-key failure
        BEGIN
            INSERT INTO tags(name, type, count) VALUES (key, Atype, Acount);
            RETURN FALSE;
        EXCEPTION WHEN unique_violation THEN
            -- Do nothing, and loop to try the UPDATE again.
        END;
    END LOOP;
END;
$$
LANGUAGE plpgsql;
    ''')
    conn.commit()

if __name__ == '__main__':
    index = 0
    logging.debug(str(sys.argv))
    if len(sys.argv) > 1:
        cmd = sys.argv[1].strip().lower()
        if cmd == 'init':
            import pwd
            if pwd.getpwuid(os.getuid()).pw_name != 'postgres':
                logging.debug("This user doesn't seem to have owner privileges, but will try anyway")
            init_db()
            logging.debug('Initialized database')
            sys.exit()
        else:
            if cmd == 'parse':
                if len(sys.argv) < 3:
                    raise Exception('No URL to parse')
                count = parse.do_parse_page(sys.argv[2], print_tag)
                print 'Parsed %d tags' % (count)
                sys.exit()
            else:
                try:
                    logging.debug('Starting from index %s' % sys.argv[1].strip())
                    index = int(sys.argv[1].strip())
                except ValueError:
                    pass
            
    conn = psycopg2.connect(database='isidabot', user='isidabot', host='localhost', password='dassad', port='5432')
    cur = conn.cursor()
    psycopg2.extensions.register_type(psycopg2.extensions.UNICODE, cur)
    try:
        while True:
            count = parse.do_parse_page('http://gelbooru.com/index.php?page=tags&s=list&pid=%d' % (index), write_tag)
            print 'Parsed %d tags' % (count)
            if not count > 0:
                sys.exit()
            if count <> 50:
                logging.debug('Invalid tags count')
            index = index + 50
            conn.commit()
            time.sleep(5)
    except Exception:
        logging.exception("Something awful happened!")
        conn.rollback()
    conn.commit()
    cur.close()
        